package chapter_4.matrix;

import java.util.Arrays;

/**
 * (Strassen). 
 *
 * 1. 将 A, B, C 分级为 n/2 X n/2 的子矩阵。
 * 2. 创建 10 个 n/2 X n/2 的矩阵 S1, S2, ..., S10.
 * 3. 用步骤 1 与步骤 2 的矩阵，递归地计算 7 个矩阵积 P1, P2, ..., P7.
 * 4. 通过 Pi 矩阵的不同组合进行加减运算，计算出结果矩阵 C 的子矩阵 C11, C12, C21, C22.
 *
 * @author 汪文波(Wang Wenbo) wenboit@163.com
 * @notes Create on 2021-12-19 10:45
 */
public class Strassen {

    public static void main(String[] args) {
        int A[][] = new int[][]{
                {1, 3, 5, 9},
                {4, 9, 3, 12},
                {9, 78, 12, 32},
                {9, 21, 41, 1}
        };

        int B[][] = new int[][]{
                {4, 9, 10, 8},
                {3, 7, 22, 15},
                {9, 13, 26, 22},
                {13, 28, 33, 19}
        };

        int[][] C = multi(A, B);

        System.out.println(Arrays.deepToString(C));
    }

    public static int[][] multi(int[][] A, int[][] B) {

        int rows = A.length;
        int[][] C = new int[rows][rows];

        if (rows == 1) {
            C[0][0] = A[0][0] * B[0][0];
            return C;
        }

        int mid = rows / 2;
        if (mid > 0) {
            int[][] A11, A12, A21, A22;
            int[][] B11, B12, B21, B22;
            int[][] C11, C12, C21, C22;
            int[][] S1, S2, S3, S4, S5, S6, S7, S8, S9, S10;
            int[][] P1, P2, P3, P4, P5, P6, P7;

            int[][][] a = partition(A, mid);
            A11 = a[0];
            A12 = a[1];
            A21 = a[2];
            A22 = a[3];

            int[][][] b = partition(B, mid);
            B11 = b[0];
            B12 = b[1];
            B21 = b[2];
            B22 = b[3];

            S1 = minus(B12, B22);
            S2 = plus(A11, A12);
            S3 = plus(A21, A22);
            S4 = minus(B21, B11);
            S5 = plus(A11, A22);
            S6 = plus(B11, B22);
            S7 = minus(A12, A22);
            S8 = plus(B21, B22);
            S9 = minus(A11, A21);
            S10 = plus(B11, B12);

            P1 = multi(A11, S1);
            P2 = multi(S2, B22);
            P3 = multi(S3, B11);
            P4 = multi(A22, S4);
            P5 = multi(S5, S6);
            P6 = multi(S7, S8);
            P7 = multi(S9, S10);

            C11 = plus(minus(plus(P5, P4), P2), P6);
            C12 = plus(P1, P2);
            C21 = plus(P3, P4);
            C22 = minus(minus(plus(P5, P1), P3), P7);

            C = combine(C11, C12, C21, C22);
        }

        return C;

    }

    private static int[][][] partition(int[][] matrix, int n) {

        int[][][] arr = new int[4][][];

        for (int i = 0; i < 4; i++) {
            arr[i] = new int[n][n];
        }

        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                arr[0][i][j] = matrix[i][j];
                arr[1][i][j] = matrix[i][j + n];
                arr[2][i][j] = matrix[i + n][j];
                arr[3][i][j] = matrix[i + n][j + n];
            }
        }

        return arr;
    }

    private static int[][] plus(int[][] mat_1, int[][] mat_2) {
        int n = mat_1.length;
        int[][] result = new int[n][n];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                result[i][j] = mat_1[i][j] + mat_2[i][j];
            }
        }
        return result;
    }

    private static int[][] minus(int[][] mat_1, int[][] mat_2) {
        int n = mat_1.length;
        int[][] result = new int[n][n];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                result[i][j] = mat_1[i][j] - mat_2[i][j];
            }
        }
        return result;
    }

    private static int[][] combine(int[][] m11, int[][] m12, int[][] m21, int[][] m22) {
        int mid = m11.length;
        int[][] c = new int[mid * 2][mid * 2];
        for (int i = 0; i < mid; i++) {
            for (int j = 0; j < mid; j++) {
                c[i][j] = m11[i][j];
                c[i][j + mid] = m12[i][j];
                c[i + mid][j] = m21[i][j];
                c[i + mid][j + mid] = m22[i][j];
            }
        }

        return c;
    }
}
